Repository navigation
Cached treetwister for twist! - #599
Conversation
f16d796 to
7584f98
Compare
|
Before I start reviewing, since twist only depends on the uncoupled indices of the trees, it is just a constant scalar factor (or thus a UniformScaling) for every FusionBlock. Id that being used here? |
|
It is, but it's slightly subtle: the code uses |
7584f98 to
fed8260
Compare
…ernel `twist!` looped over all fusion trees and scaled every subblock separately, with a cost that grows with the number of subblocks rather than the data size, and one kernel launch per subblock for GPU tensors (#595). - Add `twist(src, inds; inv)` for fusion tree pairs and blocks, returning the scalar twist factor, with `inds` given as indices or as a `Bool` mask. - Add the cached `treetwister`, a `UniqueTreeTransformer` for all sectors since a twist maps every tree pair onto itself, keyed on a fixed-length mask. - `twist!` and `twist` on a `TensorMap` go through `add_transform!` (in place for `twist!`, a single pass for `twist`); other tensor types keep the loop. - Move the shortcut for trivial permutations from `add_transform!` into `_braid!` and `_transpose!`, so `add_transform!` always applies the transformer. - Cache a device `treetwister` in the GPUArrays extension. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Adjoint destinations and sources are unwrapped as in the other index manipulations: the twist mask of an adjoint destination is relabeled onto the legs of its parent with the inverse twist, and `treetwister` takes a `conjsrc` flag, with the destination space and the permutation of the source legs following from the source space. An in-place twist of an adjoint therefore uses the same transformer as the equivalent twist of its parent. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… α, β, ...)` `twist!(t, inds)` and `twist(t, inds)` are now thin wrappers around the three-argument form, which goes through `unwrap_adjoints` and a `_twist!` kernel entry like `_braid!`. Adds `twist_pb` with Mooncake and Enzyme rules, and backend/allocator keywords for `twist`. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ist!` on Julia 1.10 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…h on Julia 1.10 Enzyme aborts in `fixup_callconv!` when compiling the wrapper while it returned the result of the three-argument `twist!`, which also broke the fermionic contraction rules. The in-place Enzyme rules call the wrapper again. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
GPU diagonal sources go through a cached `diagonal_treetwister`, analogous to `_braid!` and `_transpose!` for diagonal tensors. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
fed8260 to
29f5a17
Compare
| ) | ||
| return nothing | ||
| end | ||
| # disambiguation with the generic method |
There was a problem hiding this comment.
Is the method below only there to avoid ambiguities because of the unspecialized inds argument in the method above. Is it ever called? We do not accept a twist(!) call with inds::Index2Tuple right? And the lowering that is proposed by the invoke call doesn't work I would think, since an inds::Index2Tuple would not work in the method above?
Couldn't we just specialize the method above to the inds that we do accept/support, namely Union{Int,IndexTuple}?
| The twist is diagonal in the fusion tree basis, so this returns `src` itself together with the | ||
| scalar factor `θ`, which is shared by all trees in a block. If `inv = true`, use the inverse twist. | ||
| """ | ||
| function twist(src::Union{FusionTreePair{I}, FusionTreeBlock{I}}, inds; inv::Bool = false) where {I} |
There was a problem hiding this comment.
I see why you want to accept an inds as a mask, i.e. of type NTulpe{N, Bool}, because of how twist is implemented at the tensor level, namely by reusing some of the functionality of the levels argument of braid. But it is somewhat asymmetric that we accept a mask argument for inds in the fusiontree method, but not in the associated tensor method.
| # and store that instead of repeating the twist in the pullback each time. | ||
| # effectively, we replace `twist!` by `add ∘ twist`. | ||
| Ap = if !isa(α, Const) | ||
| Ap = twist(A.val, inds.val; inv, copy = true) |
There was a problem hiding this comment.
Should this twist call have backend and allocator passed on as keyword arguments?
| # and store that instead of repeating the twist in the pullback each time. | ||
| # effectively, we replace `twist!` by `add ∘ twist`. | ||
| Ap = if _needs_tangent(α) | ||
| Ap = twist(A, inds; inv, copy = true) |
Addresses the standalone
twist!part of #595.twist!looped overfusiontrees(t)and calledscale!on every subblock, so its cost grew with the number of subblocks rather than with the data size, and a twist of a GPU tensor launched one kernel per subblock. This PR routestwist!/twistthrough the existing index manipulation kernel with a cached tree transformer.Changes
twist(src::Union{FusionTreePair, FusionTreeBlock}, inds; inv = false) -> (src, θ): fusion-tree level twist, next topermute/braid, withindsgiven as indices or as anNTuple{N, Bool}mask (indices are converted to a fixed-length mask, so all forms are inferable).twist!(tdst, tsrc, inds, α = 1, β = 0, [backend], [allocator]; inv = false)computingtdst = β * tdst + α * twist(tsrc, inds; inv), laid out likebraid!: it unwraps adjoints throughunwrap_adjointsand calls a_twist!kernel entry mirroring_braid!.twist!(t, inds)andtwist(t, inds)are thin wrappers around it, andtwistnow also acceptsbackend/allocatorkeywords.treetwister(A, Vdst, Vsrc, p, conjsrc, twistmask, inv), with the same arguments astreebraider(the mask taking the place of the levels): always aUniqueTreeTransformer, also for non-abelian sectors, since a twist maps every tree pair onto itself with a scalar factor (no recoupling matrices / pack-matmul-unpack needed).twist!(t, inds)runs throughadd_transform!(aliasing is allowed for trivial permutations in TensorOperations), andtwistis a single out-of-place pass instead ofcopy!followed bytwist!.TensorMap/AdjointTensorMapdestinations and sources go through the twister. The twist mask is relabeled like the levels ofbraid!, and an adjoint destination flipsinv(just as it conjugatesα).twist_pbpullback with Mooncake and Enzyme rules for the three-argumenttwist!, plus Enzyme rules for the in-placetwist!(t, inds). The existing ChainRulestwistrrule already forwards keywords and is now also tested withinvandbackend/allocator.add_transform!into_braid!/_transpose!(which also skips the transformer lookup there), soadd_transform!always applies its transformer.treetwister, registered and sized like the other device caches, so a GPUtwist!is a single batched kernel launch.Benchmarks (CPU, single thread)
fℤ₂ × U(1) tensors
Vχ ⊗ VD ⊗ VD' ← Vχ ⊗ VD ⊗ VD', twisted on indices(2, 3, 5, 6), minimum of 20 runs:twist!beforetwist!aftercopyBuilding the twister on a cache miss costs ~3× a cached call in the smallest case. There is still a per-subblock cost from the individual
tensoradd!calls on the CPU, which is left for a follow-up. GPU timings were not measured.Testing
tensors/indexmanipulations.jl,tensors/contractions.jl,chainrules/linalg.jl,mooncake/indexmanipulations.jl,enzyme/flip.jlandcuda/tensors.jl(RTX A6000) pass.test/cuda/tensors.jlcomparingtwist/twist!onCuTensorMaps (and their adjoints) against the host results, and checking that the cached device transformer is used.test/tensors/indexmanipulations.jl, covering adjoint sources, adjoint destinations and in-place twists of adjoints, plus α/β accumulation.Follow-up
Folding the twist into the permutations that precede a contraction (the second part of #595) will be done on top of this, through a cached contraction plan holding the transformers.
🤖 Generated with Claude Code